import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split, GridSearchCV, StratifiedKFold
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score,
    f1_score, roc_auc_score, confusion_matrix
)

# ========== 1. 数据读取 ==========
df = pd.read_csv(r"D:\resources\data\jumbo_2020-2021M_0.25deg_ST_SSS_DO_CHL_SSH_MLD_each_year_tertile.csv")

# ========== 2. 特征与标签 ==========
target_col = "Label"
exclude_cols = ['Catch', 'Effort', 'CPUE', 'Label']

feature_cols = [col for col in df.columns if col not in exclude_cols]

# feature_cols = [
#     'ST_0', 'DO_0', 'Year', 'ST_50', 'MLD', 'Month',
#     'SSS_100', 'SSH', 'CHL', 'ST_200', 'ST_150',
#     'Lon', 'ST_500', 'DO_50', 'SSS_150'
# ]

X = df[feature_cols].copy()
y = df[target_col].astype(int)

# ========== 3. 数据划分 ==========
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

# ========== 4. 定义随机森林模型 ==========
rf = RandomForestClassifier(random_state=42, n_jobs=-1)

# ========== 5. 定义参数网格 ==========
param_grid = {
    "n_estimators": [100, 200, 300],
    "max_depth": [5, 10, 15, None],
    "min_samples_split": [2, 5, 10],
    "min_samples_leaf": [1, 2, 4],
    "max_features": ["sqrt", "log2"]
}

# ========== 6. 设置10折交叉验证 ==========
cv = StratifiedKFold(n_splits=10, shuffle=True, random_state=42)

# ========== 7. 网格搜索 ==========
grid_search = GridSearchCV(
    estimator=rf,
    param_grid=param_grid,
    # scoring="f1",
    scoring="roc_auc",
    cv=cv,
    n_jobs=-1,
    verbose=2
)

grid_search.fit(X_train, y_train)

# ========== 8. 输出最优参数与交叉验证得分 ==========
print("\n✅ 最优参数组合:")
print(grid_search.best_params_)

print(f"\n🌟 最优平均F1-score（10折）: {grid_search.best_score_:.4f}")

# ========== 9. 最优模型在测试集上评估 ==========
best_model = grid_search.best_estimator_

y_pred = best_model.predict(X_test)
y_proba = best_model.predict_proba(X_test)[:, 1]

acc = accuracy_score(y_test, y_pred)
prec = precision_score(y_test, y_pred)
rec = recall_score(y_test, y_pred)
f1 = f1_score(y_test, y_pred)
auc = roc_auc_score(y_test, y_proba)

print("\n🎯 最优随机森林模型测试集评估结果")
print(f"Accuracy  = {acc:.4f}")
print(f"Precision = {prec:.4f}")
print(f"Recall    = {rec:.4f}")
print(f"F1-score  = {f1:.4f}")
print(f"ROC-AUC   = {auc:.4f}")

print("\n混淆矩阵：")
print(confusion_matrix(y_test, y_pred))
